import os
import re
import pandas as pd
import nibabel as nib
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from skimage.transform import resize

class MedicalDataset(Dataset):
    def __init__(self, img_dir, roi_dir, label_file, target_shape=(8, 224, 224)):
        """
        参数:
            img_dir: 影像文件目录
            roi_dir: ROI文件目录
            label_file: 包含标签的Excel文件路径
            target_shape: 调整后的目标尺寸(深度, 高度, 宽度)
        """
        self.img_dir = img_dir
        self.roi_dir = roi_dir
        self.target_shape = target_shape  # (D, H, W)
        
        # 读取标签数据
        self.labels_df = pd.read_excel(label_file, dtype={'p_id': str})  # 强制p_id列为字符串
        self.label_map = dict(zip(self.labels_df['p_id'], self.labels_df['label']))
        
        # 构建样本列表
        self.samples = []
        pattern = re.compile(r'(\d+)_T2_(axi|cor|sag)_\d+\.nii\.gz')
        
        for img_file in os.listdir(img_dir):
            match = pattern.match(img_file)
            if match:
                p_id, scan_type = match.groups()
                p_id = str(p_id)  # 确保类型一致
                roi_file = f"{p_id}_T2_{scan_type}_roi.nii.gz"
                roi_path = os.path.join(roi_dir, roi_file)
                
                if p_id in self.label_map and os.path.exists(roi_path):
                    self.samples.append({
                        'img': os.path.join(img_dir, img_file),
                        'roi': roi_path,
                        'label': self.label_map[p_id],
                        'scan_type': scan_type
                    })

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        
        # 加载数据
        img = nib.load(sample['img']).get_fdata()
        roi = nib.load(sample['roi']).get_fdata()
        
        # 统一轴向 (假设原始数据为 RAS 坐标系)
        img = np.transpose(img, (2, 0, 1))  # 调整为 (D, H, W)
        roi = np.transpose(roi, (2, 0, 1))
        
        # 调整尺寸
        img = self._process_volume(img)
        roi = self._process_volume(roi, is_mask=True)
        
        # 应用ROI
        masked_img = img * (roi > 0.5)  # 假设ROI为概率图
        masked_img = self._normalize(masked_img)
        
        # 转换为张量
        tensor_img = torch.from_numpy(masked_img).float().unsqueeze(0)  # (C, D, H, W)
        label = torch.tensor(sample['label'], dtype=torch.long)
        
        return tensor_img, label

    def _process_volume(self, volume, is_mask=False):
        """处理3D体积数据"""
        # 调整尺寸
        if volume.shape != self.target_shape:
            order = 0 if is_mask else 3  # mask用最近邻插值，图像用三次样条
            volume = resize(volume, 
                           self.target_shape,
                           order=order,
                           mode='constant',
                           preserve_range=True,
                           anti_aliasing=False)
        return volume

    def _normalize(self, volume):
        """自适应直方图归一化"""
        # 仅使用ROI区域计算统计量
        non_zero = volume[volume > 0]
        if len(non_zero) == 0:
            return volume
            
        mean = np.mean(non_zero)
        std = np.std(non_zero)
        return (volume - mean) / (std + 1e-8)

# 使用示例
if __name__ == "__main__":
    # 初始化数据集
    dataset = MedicalDataset(
        img_dir='./datasets/three_angle/train/img',
        roi_dir='./datasets/three_angle/train/roi',
        label_file='./datasets/three_angle/label.xlsx'
    )
    
    # 创建数据加载器
    dataloader = DataLoader(dataset, 
                           batch_size=2,
                           shuffle=True,
                           pin_memory=True)
    
    # 验证数据形状
    for data, labels in dataloader:
        print(f"数据形状: {data.shape}")  # 应输出 torch.Size([batch, 1, 8, 224, 224])
        print(f"标签形状: {labels.shape}")
        break